package net.cyue.jpa;

import jakarta.persistence.EntityManager;
import jakarta.persistence.NoResultException;
import jakarta.persistence.Query;
import jakarta.persistence.criteria.CriteriaQuery;
import jakarta.persistence.criteria.Path;
import jakarta.persistence.criteria.Root;

import java.util.List;
import java.util.Map;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.atomic.AtomicReference;
import java.util.function.Consumer;

@FunctionalInterface
interface Function {
    void invoke();
}

public class EntityManagerWrapper {
    private final EntityManager entityManager;
    private boolean using = false;

    public EntityManagerWrapper(EntityManager entityManager) {
        this.entityManager = entityManager;
    }

    public void close() {
        this.entityManager.close();
    }
    public boolean isUsing() {
        return this.using;
    }

    public Query createQuery(String jpql, Object... parameters) {
        Query q = this.entityManager.createQuery(jpql);
        for (int i = 1; i <= parameters.length; i++) {
            q.setParameter(i, parameters[i - 1]);
        }
        return q;
    }
    public Query createQuery(String jpql, Map<String, ?> parameters) {
        Query q = this.entityManager.createQuery(jpql);
        for (String key : parameters.keySet()) {
            q.setParameter(key, parameters.get(key));
        }
        return q;
    }
    public <T> Query createQuery(Class<T> clz, String... columnName) {
        CriteriaQuery<T> cq = this.entityManager.getCriteriaBuilder().createQuery(clz);
        Root<T> root = cq.from(clz);
        Path<?>[] paths = new Path[columnName.length];
        for (int i = 0; i < columnName.length; i++) {
            paths[i] = root.get(columnName[i]);
        }
        CriteriaQuery<T> select = cq.multiselect(paths);
        return this.entityManager.createQuery(select);
    }

    private void use(Function fn) {
        if (this.using) {
            throw new RuntimeException("EntityManager is already using");
        }
        this.using = true;
        fn.invoke();
        this.using = false;
    }
    public void transaction(Consumer<EntityManager> work) {
        this.use(
            () -> {
                final var transaction = this.entityManager.getTransaction();
                try {
                    transaction.begin();
                    work.accept(this.entityManager);
                    transaction.commit();
                } catch (Exception e) {
                    transaction.rollback();
                    throw e;
                }
            }
        );
    }
    public void transactionWrapper(Consumer<EntityManagerWrapper> work) {
        this.use(
            () -> {
                final var transaction = this.entityManager.getTransaction();
                try {
                    transaction.begin();
                    work.accept(this);
                    transaction.commit();
                } catch (Exception e) {
                    transaction.rollback();
                    throw e;
                }
            }
        );
    }

    public void query(Consumer<Query> work, String jpql, Object... parameters) {
        if (
            jpql == null ||
            jpql.trim().isEmpty()
        ) {
            throw new IllegalArgumentException("jpql is empty");
        }

        this.transaction((manager) -> {
            Query q = this.createQuery(jpql, parameters);
            work.accept(q);
        });
    }
    public void query(Consumer<Query> work, String jpql, Map<String, ?> parameters) {
        if (
            jpql == null ||
            jpql.trim().isEmpty()
        ) {
            throw new IllegalArgumentException("jpql is empty");
        }

        this.transaction((manager) -> {
            Query q = this.createQuery(jpql, parameters);
            work.accept(q);
        });
    }
    public <T> void query(Consumer<Query> work, Class<T> clz, String... columnName) {
        if (clz == null) {
            return;
        }

        this.transaction((manager) -> {
            Query q = this.createQuery(clz, columnName);
            work.accept(q);
        });
    }

    public <T> List<T> select(Class<T> clz, String... columnName) {
        AtomicReference<List<T>> ref = new AtomicReference<>();
        this.query(
            (q) -> {
                ref.set(q.getResultList());
            },
            clz,
            columnName
        );
        return ref.get();
    }
    public <T> T selectOne(Class<T> clz, String... columnName) {
        AtomicReference<T> ref = new AtomicReference<>();
        this.query(
            (q) -> {
                ref.set((T) q.getSingleResult());
            },
            clz,
            columnName
        );
        return ref.get();
    }
    public List select(String jpql, Object... parameters) {
        AtomicReference<List<?>> ref = new AtomicReference<>();
        this.query(
            (q) -> {
                ref.set(q.getResultList());
            },
            jpql,
            parameters
        );
        return ref.get();
    }
    public List select(String jpql, Map<String, ?> parameters) {
        AtomicReference<List<?>> ref = new AtomicReference<>();
        this.query(
            (q) -> {
                ref.set(q.getResultList());
            },
            jpql,
            parameters
        );
        return ref.get();
    }
    public Object selectOne(String jpql, Object... parameters) {
        AtomicReference<Object> ref = new AtomicReference<>();
        this.query(
            (q) -> {
                try {
                    ref.set(q.getSingleResult());
                } catch (NoResultException e) {
                    ref.set(null);
                }
            },
            jpql,
            parameters
        );
        return ref.get();
    }
    public Object selectOne(String jpql, Map<String, ?> parameters) {
        AtomicReference<Object> ref = new AtomicReference<>();
        this.query(
            (q) -> {
                try {
                    ref.set(q.getSingleResult());
                } catch (NoResultException e) {
                    ref.set(null);
                }
            },
            jpql,
            parameters
        );
        return ref.get();
    }

    public int executeUpdate(String jpql, Object... parameters) {
        AtomicInteger ref = new AtomicInteger();
        this.query(
            (q) -> {
                ref.set(q.executeUpdate());
            },
            jpql,
            parameters
        );
        return ref.get();
    }
    public int executeUpdate(String jpql, Map<String, ?> parameters) {
        AtomicInteger ref = new AtomicInteger();
        this.query(
            (q) -> {
                ref.set(q.executeUpdate());
            },
            jpql,
            parameters
        );
        return ref.get();
    }


    public void persist(Object obj) {
        if (obj == null) {
            return;
        }
        this.transaction((manager) -> manager.persist(obj));
    }

    public <T> T find(Class<T> clz, int id) {
        if (clz == null) {
            return null;
        }
        if (id < 0) {
            return null;
        }

        AtomicReference<T> ref = new AtomicReference<>();
        this.transaction((manager) -> {
            ref.set(manager.find(clz, id));
        });
        return ref.get();
    }

    public void remove(Object obj) {
        if (obj == null) {
            return;
        }
        this.transaction((manager) -> manager.remove(obj));
    }
    public <T> void remove(Class<T> clz, int id) {
        T obj = this.find(clz, id);
        this.remove(obj);
    }
}
